import numpy as np
import random
from PIL import Image
import cv2
from util_flux import process_img_1024

def load_yolo():
    from ultralytics import YOLOE
    model = YOLOE("yoloe-11l-seg.pt")  # or select yoloe-11s/m-seg.pt for different sizes

    # Set text prompt to detect person and bus. You only need to do this once after you load the model.
    names = ["clothing",'garment','UpperClothing','LowerClothing','FullbodyClothing']
    model.set_classes(names, model.get_text_pe(names))
    return model

def get_mask_by_yolo(model,img_path = '' , img_pil = None):
    img_pil = img_pil if img_pil is not None else Image.open(img_path)
    results = model.predict(img_pil)

    res = results[0]

    # pdb.set_trace()
    # 假设 mask 是形状为 [1, H, W] 的 torch.Tensor
    try:
        mask = res.masks[0].data.squeeze(0).cpu().numpy()  # 去除批次维度并转为 numpy 数组

        # 将 mask 二值化（0-1 转为 0-255）
        mask_uint8 = (mask * 255).astype(np.uint8)  # mask !
        # mask_pil
        mask_pil = Image.fromarray(mask_uint8)
        mask_pil = process_img_1024( '' , img_pil=mask_pil,
                        target_shape=img_pil.size , pad_color=(0,0,0) )
        mask_pil = mask_pil.convert('L')
        mask_np = np.array( mask_pil )
    except:
        mask_pil = None
        mask_np = None

    return img_pil , mask_np , mask_pil


def center_crop(image, size=224):
    # min_size = min(control_size)
    # crop_ratio = min_size//size + 1
    # size *= crop_ratio
    w, h = image.size
    # min_size = min(w,h)//2
    # size = min(w, h, min_size)
    x = w//4
    y = h//4

    return image.crop((x,y,x+w//2,y+h//2))

def erode_to_target_area(mask, target_ratio, max_iter=1000, kernel_size=3):
    """
    通过形态学腐蚀精确控制面积缩减
    :param mask: 输入二值掩膜(0/255)
    :param target_ratio: 目标面积比例(0-1)
    :param max_iter: 最大迭代次数
    :param kernel_size: 腐蚀核大小(建议3或5)
    :return: 侵蚀后的掩膜
    """
    # 初始化
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
    current_mask = mask.copy()
    original_area = np.sum(current_mask > 0)
    target_area = int(original_area * target_ratio)
    
    # 逐步腐蚀
    for i in range(max_iter):
        eroded = cv2.erode(current_mask, kernel)
        current_area = np.sum(eroded > 0)
        
        # 终止条件
        if current_area <= target_area or current_area == 0:
            break
            
        current_mask = eroded
    
    # 最终面积校验
    final_ratio = np.sum(current_mask > 0) / original_area
    print(f"目标比例: {target_ratio:.0%} → 实际比例: {final_ratio:.1%} (迭代次数: {i})")
    
    return current_mask

def get_erosed_mask_by_radtio(mask,ratio):
# for ratio in [0.9, 0.7, 0.5, 0.3,0.1]:
    scaled = erode_to_target_area(mask.copy(), ratio)
    # print( np.sum(scaled > 0) / np.sum(mask > 0))
    # cv2.imwrite(f'mask_scaled_{ratio:.1f}.png', scaled)
    scaled_pil = Image.fromarray(scaled)
    return scaled,scaled_pil


def add_random_holes(mask_pil, min_holes=1,max_holes=5, hole_size_range=(10, 50)):
    mask = np.array(mask_pil) # 读取单通道Mask
    if len(mask.shape) == 3:  # 如果是RGB图像，转换为单通道
        mask = mask[:, :, 0]
    mask = (mask > 0).astype(np.uint8)  # 二值化
    """
    在现有Mask上随机挖洞
    :param mask: 二值化Mask数组（0表示背景，1表示目标）
    :param max_holes: 最大挖洞数量
    :param hole_size_range: 洞的尺寸范围(最小像素, 最大像素)
    :return: 带洞的Mask
    """
    mask = mask.copy()
    h, w = mask.shape
    n_holes = random.randint(min_holes, max_holes)
    
    for _ in range(n_holes):
        # 随机生成洞的尺寸
        hole_w = random.randint(*hole_size_range)
        hole_h = random.randint(*hole_size_range)
        
        # 只在Mask有效区域内挖洞
        ys, xs = np.where(mask == 1)
        if len(ys) == 0:
            break
            
        # 随机选择中心点
        idx = random.choice(range(len(ys)))
        center_x, center_y = xs[idx], ys[idx]
        
        # 计算洞的边界（防止越界）
        x1 = max(0, center_x - hole_w//2)
        y1 = max(0, center_y - hole_h//2)
        x2 = min(w, x1 + hole_w)
        y2 = min(h, y1 + hole_h)
        
        # 挖洞（设置为背景值0）
        mask[y1:y2, x1:x2] = 0
        
    mask_pil = Image.fromarray(mask*255)
    return mask,mask_pil

# 使用示例
# mask = np.array(Image.open("mask.png").convert("L"))  # 读取单通道Mask
# mask = (mask > 0).astype(np.uint8)  # 二值化
# hole_mask = add_random_holes(mask, max_holes=8, hole_size_range=(15, 60))
# Image.fromarray(hole_mask * 255).save("hole_mask.png")